#这个文件用于提取部分类别图片，主要分为三部分
#1. 首先选取只含class_name_path文件的部分类别
#2. 将对应类别图片放入img_savepath
#3. 将对应类别xml写入ann_savepath
#4. 将对应类别xml写入trainval_path文件的txt里
import os
import shutil
from tqdm import tqdm
import shutil
import cv2
import random
from PIL import Image, ImageDraw
import xml.etree.ElementTree as ET

ann_path_origin = r'/home/xys/datasets/dataset_che/dataset-larger/Annotations/'        #要提取的ann标签文件路径
img_path_origin = r'/home/xys/datasets/dataset_che/dataset-larger/imgs/'         #要提取的Image文件路径
ann_savepath = r'/home/xys/datasets/dataset_che/dataset-larger/ann/'        #要保存到的目标annotation路径
img_savepath = r'/home/xys/datasets/dataset_che/dataset-larger/JPEGImages/'         #要保存到的Images文件路径
class_name_path = r'/home/xys/datasets/dataset_che/dataset-larger/classes.txt'      #存储的类别信息,用来提取
trainval_path = r'/home/xys/datasets/dataset_che/dataset-larger/trainval.txt'

#-----------------------函数定义------------------------------------
def get_class(path):
    with open(path) as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    return class_names

def mkr(path):
    '''
    如果path存在，就先删除所有内容再创建，否则直接创建
    '''
    if os.path.exists(path):#先删除再创建
        shutil.rmtree(path)#递归删除目录树
        os.mkdir(path)#创建目录
    else:
        os.mkdir(path)#如果不存在则直接创建

def write_xml(anno_path, head, objs, tail):
    '''
    写入xml文件
    anno_path:路径
    head:头部
    objs:
    tail:尾部
    '''
    f = open(anno_path, "w")#以w打开文件流
    f.write(head)#写入头部
    for obj in objs:#将该objs所有内容写入
        f.write(objstr % (obj[0], obj[1], obj[2], obj[3], obj[4]))
    f.write(tail)#写入尾部
    f.close()

names = locals()
classes_name = 'vehicle'                    #读取类别信息
mkr(ann_savepath)                               #创建文件夹
mkr(img_savepath)

headstr = """\
<annotation>
    <folder>VOC</folder>
    <filename>%s</filename>
    <source>
        <database>My Database</database>
        <annotation>COCO</annotation>
        <image>flickr</image>
        <flickrid>NULL</flickrid>
    </source>
    <owner>
        <flickrid>NULL</flickrid>
        <name>company</name>
    </owner>
    <size>
        <width>%d</width>
        <height>%d</height>
        <depth>%d</depth>
    </size>
    <segmented>0</segmented>
"""
objstr = """\
    <object>
        <name>%s</name>
        <pose>Unspecified</pose>
        <truncated>0</truncated>
        <difficult>0</difficult>
        <bndbox>
            <xmin>%d</xmin>
            <ymin>%d</ymin>
            <xmax>%d</xmax>
            <ymax>%d</ymax>
        </bndbox>
    </object>
"""

tailstr = '''\
</annotation>
'''

dir = os.listdir(ann_path_origin)
random.shuffle(dir)
for xml_file_name in tqdm(dir,ncols=150):
    boundingboxes = []
    objs = []                                           #有这个类，将这个类的所有boundingbox和类别信息写入
    file_name = xml_file_name[:-4]
    xml_file = open(ann_path_origin+xml_file_name)      #打开xml文件
    tree = ET.parse(xml_file)
    root = tree.getroot()                               #得到根
    for obj in root.iter('object'):
        diffcult = obj.find('difficult').text         #找到difficult，难以识别的类
        cls_name = obj.find('name').text
        if cls_name not in classes_name or int(diffcult) == 1:
            continue

        #--------------------下面的语句表示存在这个boundingbox，将此boundingbox的坐标类别写入到xml里-------------------
        have = False if cls_name not in classes_name else True
        xmlbox = obj.find('bndbox')                 #找到boundingbox标签
        xmin = int(float(xmlbox.find('xmin').text))
        ymin = int(float(xmlbox.find('ymin').text))
        xmax = int(float(xmlbox.find('xmax').text))
        ymax = int(float(xmlbox.find('ymax').text))
        b = [xmin,ymin,xmax,ymax]

        #------------存入boundingbox的列表，用来放入txt里---------------------
        cls_id = classes_name.index(cls_name)
        str_bb =",".join([str(a) for a in b])+','+str(cls_id)
        boundingboxes.append(str_bb)

        #-------------------存入boundingbox的坐标数据，用来放入xml里------------------
        obj = [cls_name, xmin, ymin, xmax, ymax]  # 将类别名，boundingbox的坐标存入objs[]
        objs.append(obj)

    if len(boundingboxes):                          #如果这张图片有我们需要的类，则保存这张图片,并写入到txt文件里
        shutil.copy(img_path_origin+file_name+".jpg",img_savepath+file_name+".jpg")
        str_path = img_savepath+file_name+".jpg "+" ".join([str(b) for b in boundingboxes])+'\n'
        txt_file = open(trainval_path,'a')    #打开txt文件流
        txt_file.write(str_path)                    #写入txt

        img = cv2.imread(img_path_origin+file_name+".jpg")  # 打开/读取图片
        head = headstr % (file_name, img.shape[1], img.shape[0], img.shape[2])
        write_xml(ann_savepath+xml_file_name,head=head,tail = tailstr,objs=objs)

        txt_file.close()
    xml_file.close()
print("---------结束-----------------")
